from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
import os
import psutil
from datetime import datetime
import json
from collections import deque
import time
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.utils import formataddr
import logging
from flask_sock import Sock
import gzip
import redis

class User(UserMixin):
    def __init__(self, username):
        self.id = username

class WebInterface:
    def __init__(self, defender):
        self.defender = defender
        
        # 创建Flask应用
        self.app = Flask(__name__)
        self.app.secret_key = os.urandom(24)
        
        # 设置日志
        self.setup_logging()
        
        # 初始化Flask-Login
        self.login_manager = LoginManager()
        self.login_manager.init_app(self.app)
        self.login_manager.login_view = 'login'
        
        # 初始化 WebSocket
        self.sock = Sock(self.app)
        
        # 初始化 Redis
        self.redis = redis.Redis(
            host='localhost',
            port=6379,
            db=0,
            decode_responses=True
        )
        
        # 添加进程列表缓存
        self.process_cache = None
        self.last_process_update = 0
        self.process_cache_ttl = 3  # 缓存3秒
        
        # 用户加载函数
        @self.login_manager.user_loader
        def load_user(username):
            return User(username)
        
        # 注册路由
        self.register_routes()
        
    def setup_logging(self):
        """设置日志"""
        logging.basicConfig(
            filename='web.log',
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        # 添加控制台输出
        console = logging.StreamHandler()
        console.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        console.setFormatter(formatter)
        logging.getLogger('').addHandler(console)
        
    def register_routes(self):
        """注册所有路由"""
        # 首页
        @self.app.route('/')
        def index():
            try:
                # 开发阶段自动登录
                if not hasattr(self, '_auto_logged_in'):
                    user = User('admin')
                    login_user(user)
                    self._auto_logged_in = True
                
                system_stats = self.defender.get_system_stats()
                if not system_stats:
                    raise Exception("无法获取系统状态")
                    
                return render_template('index.html', 
                                    banned_ips=list(self.defender.banned_ips),
                                    whitelist_ips=list(self.defender.whitelist_ips),
                                    connection_limit=self.defender.connection_limit,
                                    system_stats=system_stats,
                                    alert_settings=self.defender.alert_settings)
            except Exception as e:
                logging.error(f"渲染首页失败: {str(e)}")
                return f"服务器内部错误: {str(e)}", 500

        # 登录（开发阶段可以注释掉或保留作为备用）
        @self.app.route('/login', methods=['GET', 'POST'])
        def login():
            # 开发阶段直接登录并跳转
            user = User('admin')
            login_user(user)
            return redirect(url_for('index'))

        # 登出
        @self.app.route('/logout')
        @login_required
        def logout():
            logout_user()
            return redirect(url_for('login'))

        # API路由
        @self.app.route('/api/stats')
        def get_stats():
            """获取系统状态"""
            try:
                stats = self.defender.get_system_stats()
                if not stats:
                    raise Exception("无法获取系统状态")
                    
                # 构建响应数据
                response_data = {
                    'timestamp': stats.get('timestamp', ''),
                    'cpu': {
                        'percent': stats.get('cpu', {}).get('percent', 0),
                        'cores': stats.get('cpu', {}).get('cores', 1),
                        'frequency': stats.get('cpu', {}).get('frequency', {
                            'current': 0,
                            'min': 0,
                            'max': 0
                        }),
                        'load_average': stats.get('cpu', {}).get('load_average', {
                            '1min': 0,
                            '5min': 0,
                            '15min': 0
                        })
                    },
                    'memory': {
                        'total': stats.get('memory', {}).get('total', 0),
                        'used': stats.get('memory', {}).get('used', 0),
                        'free': stats.get('memory', {}).get('free', 0),
                        'percent': stats.get('memory', {}).get('percent', 0),
                        'swap': stats.get('memory', {}).get('swap', {
                            'total': 0,
                            'used': 0,
                            'free': 0,
                            'percent': 0
                        })
                    },
                    'disk': {
                        'total': stats.get('disk', {}).get('total', 0),
                        'used': stats.get('disk', {}).get('used', 0),
                        'free': stats.get('disk', {}).get('free', 0),
                        'percent': stats.get('disk', {}).get('percent', 0),
                        'io': stats.get('disk', {}).get('io', {
                            'read_bytes': 0,
                            'write_bytes': 0
                        })
                    },
                    'network': {
                        'bytes_sent': stats.get('network', {}).get('bytes_sent', 0),
                        'bytes_recv': stats.get('network', {}).get('bytes_recv', 0),
                        'connections': stats.get('network', {}).get('connections', {
                            'ESTABLISHED': 0,
                            'TIME_WAIT': 0,
                            'CLOSE_WAIT': 0,
                            'LISTEN': 0
                        }),
                        'interfaces': stats.get('network', {}).get('interfaces', {})
                    },
                    'processes': {
                        'total': stats.get('processes', {}).get('total', 0),
                        'top': stats.get('processes', {}).get('top', [])
                    }
                }
                
                return jsonify(response_data)
            except Exception as e:
                logging.error(f"获取状态失败: {str(e)}")
                return jsonify({
                    'error': str(e),
                    'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'cpu': {
                        'percent': 0,
                        'cores': 1,
                        'frequency': {'current': 0, 'min': 0, 'max': 0},
                        'load_average': {'1min': 0, '5min': 0, '15min': 0}
                    },
                    'memory': {
                        'total': 0,
                        'used': 0,
                        'free': 0,
                        'percent': 0,
                        'swap': {'total': 0, 'used': 0, 'free': 0, 'percent': 0}
                    },
                    'disk': {
                        'total': 0,
                        'used': 0,
                        'free': 0,
                        'percent': 0,
                        'io': {'read_bytes': 0, 'write_bytes': 0}
                    },
                    'network': {
                        'bytes_sent': 0,
                        'bytes_recv': 0,
                        'connections': {
                            'ESTABLISHED': 0,
                            'TIME_WAIT': 0,
                            'CLOSE_WAIT': 0,
                            'LISTEN': 0
                        },
                        'interfaces': {}
                    },
                    'processes': {
                        'total': 0,
                        'top': []
                    }
                }), 500

        @self.app.route('/api/ban_ip', methods=['POST'])
        def ban_ip():
            try:
                ip = request.json.get('ip')
                if not ip:
                    return jsonify({'status': 'error', 'message': '无效的IP地址'})
                self.defender.ban_ip(ip)
                return jsonify({'status': 'success'})
            except Exception as e:
                logging.error(f"封禁IP失败: {str(e)}")
                return jsonify({'status': 'error', 'message': str(e)})

        @self.app.route('/api/unban_ip', methods=['POST'])
        @login_required
        def unban_ip():
            ip = request.json.get('ip')
            if ip in self.defender.banned_ips:
                self.defender.unban_ip(ip)
                return jsonify({'status': 'success'})
            return jsonify({'status': 'error', 'message': 'IP未被封禁'})

        @self.app.route('/api/update_limit', methods=['POST'])
        @login_required
        def update_limit():
            limit = request.json.get('limit')
            try:
                limit = int(limit)
                self.defender.connection_limit = limit
                return jsonify({'status': 'success'})
            except:
                return jsonify({'status': 'error', 'message': '无效的限制值'})

        @self.app.route('/api/whitelist', methods=['GET'])
        @login_required
        def get_whitelist():
            return jsonify(list(self.defender.whitelist_ips))

        @self.app.route('/api/whitelist/add', methods=['POST'])
        @login_required
        def add_to_whitelist():
            ip = request.json.get('ip')
            if ip:
                self.defender.add_to_whitelist(ip)
                return jsonify({'status': 'success'})
            return jsonify({'status': 'error', 'message': '无效的IP地址'})

        @self.app.route('/api/whitelist/remove', methods=['POST'])
        @login_required
        def remove_from_whitelist():
            ip = request.json.get('ip')
            if ip in self.defender.whitelist_ips:
                self.defender.remove_from_whitelist(ip)
                return jsonify({'status': 'success'})
            return jsonify({'status': 'error', 'message': 'IP不在白名单中'})

        @self.app.route('/api/logs')
        @login_required
        def get_logs():
            lines = request.args.get('lines', 100, type=int)
            return jsonify(self.defender.get_recent_logs(lines))

        @self.app.route('/api/alert-settings', methods=['GET', 'POST'])
        @login_required
        def alert_settings():
            if request.method == 'POST':
                settings = request.json
                # 验证SMTP设置
                if settings.get('email_notification'):
                    try:
                        with smtplib.SMTP(settings['smtp_server'], int(settings['smtp_port'])) as server:
                            server.starttls()
                            server.login(settings['smtp_username'], settings['smtp_password'])
                            # 发送测试邮件
                            subject = "Linux防护系统 - 测试邮件"
                            message = "邮件告警功能测试成功！"
                            msg = MIMEMultipart()
                            msg['From'] = formataddr(("Linux防护系统", settings['smtp_username']))
                            msg['To'] = settings['email_address']
                            msg['Subject'] = subject
                            msg.attach(MIMEText(message))
                            server.send_message(msg)
                    except Exception as e:
                        return jsonify({'status': 'error', 'message': f'SMTP设置验证失败: {str(e)}'})
                    
                self.defender.update_alert_settings(settings)
                return jsonify({'status': 'success'})
            return jsonify(self.defender.alert_settings)

        @self.app.route('/api/history')
        @login_required
        def get_history():
            return jsonify(self.defender.get_history_data())

        @self.app.route('/api/processes')
        def get_processes():
            """获取进程列表（带分页和缓存）"""
            try:
                page = int(request.args.get('page', 1))
                size = int(request.args.get('size', 20))
                
                current_time = time.time()
                
                # 检查缓存是否有效
                if (self.process_cache is not None and 
                    current_time - self.last_process_update < self.process_cache_ttl):
                    processes = self.process_cache
                else:
                    # 获取新的进程列表，只获取必要的信息
                    processes = []
                    for proc in psutil.process_iter(['pid', 'name', 'username', 'status']):
                        try:
                            pinfo = proc.info
                            # 单独获取 CPU 和内存使用率，避免获取其他不必要的信息
                            with proc.oneshot():  # 使用 oneshot() 减少系统调用
                                pinfo['cpu_percent'] = proc.cpu_percent(interval=None)
                                pinfo['memory_percent'] = proc.memory_percent()
                            processes.append(pinfo)
                        except (psutil.NoSuchProcess, psutil.AccessDenied):
                            continue
                    
                    # 更新缓存
                    self.process_cache = processes
                    self.last_process_update = current_time
                
                # 计算进程状态统计（使用缓存的数据）
                total = len(processes)
                running = sum(1 for p in processes if p.get('status') == 'running')
                sleeping = sum(1 for p in processes if p.get('status') == 'sleeping')
                zombie = sum(1 for p in processes if p.get('status') == 'zombie')
                
                # 排序（使用缓存的数据）
                processes.sort(key=lambda x: x.get('cpu_percent', 0) or 0, reverse=True)
                
                # 分页
                start = (page - 1) * size
                end = start + size
                paginated_processes = processes[start:end]
                
                return jsonify({
                    'processes': paginated_processes,
                    'total': total,
                    'running': running,
                    'sleeping': sleeping,
                    'zombie': zombie,
                    'page': page,
                    'size': size,
                    'pages': (total + size - 1) // size
                })
                
            except Exception as e:
                logging.error(f"获取进程列表失败: {str(e)}")
                return jsonify({
                    'error': str(e),
                    'processes': [],
                    'total': 0,
                    'running': 0,
                    'sleeping': 0,
                    'zombie': 0,
                    'page': 1,
                    'size': size,
                    'pages': 0
                }), 500

        @self.app.route('/api/process/<int:pid>')
        def get_process(pid):
            """获取进程详情"""
            try:
                process = self.defender.get_process_info(pid)
                if process:
                    return jsonify(process)
                return jsonify({'error': '进程不存在'}), 404
            except Exception as e:
                logging.error(f"获取进程详情失败: {str(e)}")
                return jsonify({'error': str(e)}), 500

        @self.app.route('/api/process/<int:pid>/kill', methods=['POST'])
        def kill_process(pid):
            """终止进程"""
            try:
                if self.defender.kill_process(pid):
                    return jsonify({'status': 'success'})
                return jsonify({'status': 'error', 'message': '终止进程失败'})
            except Exception as e:
                logging.error(f"终止进程失败: {str(e)}")
                return jsonify({'error': str(e)}), 500

        # 添加 WebSocket 路由
        @self.sock.route('/ws/stats')
        def stats_socket(ws):
            """WebSocket处理函数"""
            try:
                # 检查会话状态
                if not hasattr(current_user, 'is_authenticated') or not current_user.is_authenticated:
                    ws.send(json.dumps({'error': 'unauthorized'}))
                    return
                    
                while True:
                    try:
                        # 等待客户端请求
                        message = ws.receive()
                        if message == 'get_stats':
                            # 获取最新状态
                            current_stats = self.get_cached_stats()
                            if current_stats:
                                # 发送数据
                                try:
                                    compressed = self.compress_data(current_stats)
                                    ws.send(compressed)
                                except Exception as e:
                                    logging.error(f"发送数据失败: {str(e)}")
                    except Exception as e:
                        logging.error(f"WebSocket处理消息失败: {str(e)}")
                        break
            except Exception as e:
                logging.error(f"WebSocket连接错误: {str(e)}")
            finally:
                logging.info("WebSocket连接已关闭")

        @self.app.route('/api/check-session')
        def check_session():
            """检查会话状态"""
            try:
                is_authenticated = (hasattr(current_user, 'is_authenticated') and 
                                  current_user.is_authenticated)
                return jsonify({
                    'authenticated': is_authenticated,
                    'username': current_user.id if is_authenticated else None
                })
            except Exception as e:
                logging.error(f"检查会话状态失败: {str(e)}")
                return jsonify({
                    'authenticated': False,
                    'error': str(e)
                }), 500

    def get_cached_stats(self):
        """从Redis获取缓存的系统状态"""
        try:
            # 尝试从Redis获取缓存
            cached = self.redis.get('system_stats')
            if cached:
                return json.loads(cached)
            
            # 如果没有缓存，获取新数据
            stats = self.defender.get_system_stats()
            
            # 缓存数据（5秒过期）
            self.redis.setex(
                'system_stats',
                5,  # 过期时间5秒
                json.dumps(stats)
            )
            
            return stats
        except redis.RedisError as e:
            logging.error(f"Redis错误: {str(e)}")
            return self.defender.get_system_stats()

    def calculate_stats_diff(self, old_stats, new_stats):
        """计算两次状态之间的差异"""
        diff = {}
        
        # 只比较关键指标
        for key in ['cpu', 'memory', 'disk', 'network']:
            if key in new_stats:
                if key not in old_stats or new_stats[key] != old_stats[key]:
                    diff[key] = new_stats[key]
        
        # 添加时间戳
        if diff:
            diff['timestamp'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
        return diff

    def compress_data(self, data):
        """压缩数据"""
        try:
            json_str = json.dumps(data)
            compressed = gzip.compress(json_str.encode('utf-8'))
            return compressed
        except Exception as e:
            logging.error(f"压缩数据失败: {str(e)}")
            return json.dumps(data)

    def start(self, host='0.0.0.0', port=5000):
        """启动Web服务"""
        try:
            logging.info("Web服务正在启动...")
            self.app.run(host=host, port=port, debug=False, threaded=True)
        except Exception as e:
            logging.error(f"Web服务启动失败: {str(e)}")
            raise 